# Pose绘制
import os
import cv2
import numpy as np
import mediapipe as mp
import argparse
def extract_skeleton_with_validation(input_image_path):
    # 初始化 MediaPipe Pose
    mp_pose = mp.solutions.pose
    pose = mp_pose.Pose(static_image_mode=True, min_detection_confidence=0.5)
    mp_drawing = mp.solutions.drawing_utils

    # 读取输入图像
    image = cv2.imread(input_image_path)
    if image is None:
        print("无法读取输入图像，请检查路径！")
        return

    # 将图像从 BGR 转为 RGB
    image_rgb = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)

    # 处理图像，提取姿态
    results = pose.process(image_rgb)

    # 如果检测到姿态
    if results.pose_landmarks:
        keypoints = results.pose_landmarks.landmark

        # 检查关键点是否足够
        visible_keypoints = [kp for kp in keypoints if kp.visibility > 0.5]
        if len(visible_keypoints) < 15:
            # 在图像上显示“不合格”字样
            cv2.putText(
                image, 
                "Picture Unqualified", 
                (50, 50), 
                cv2.FONT_HERSHEY_SIMPLEX, 
                1, 
                (0, 0, 255), 
                2
            )
            # 保存并退出
            print("关键点不足，图片不合格")
            return

        # 创建一个空白图像（黑色背景）
        skeleton_image = np.zeros_like(image)

        # 提取鼻尖（头部）和颈部关键点
        image_height, image_width = image.shape[:2]
        head_point = (int(keypoints[0].x * image_width), int(keypoints[0].y * image_height))  # 鼻尖
        neck_point = (
            int((keypoints[11].x + keypoints[12].x) / 2 * image_width),  # 左肩和右肩的中点作为颈部
            int((keypoints[11].y + keypoints[12].y) / 2 * image_height)
        )

        # 绘制头部点和颈部点
        cv2.circle(skeleton_image, head_point, radius=5, color=(0, 255, 0), thickness=-1)  # 头部点
        cv2.circle(skeleton_image, neck_point, radius=5, color=(0, 255, 0), thickness=-1)  # 颈部点

        # 绘制头部和颈部的连线
        cv2.line(skeleton_image, head_point, neck_point, color=(255, 0, 0), thickness=2)

        # 定义骨架的连接线
        skeleton_connections = [
            (11, 12),  # 左右肩膀
            (11, 13),  # 左肩到左肘
            (13, 15),  # 左肘到左手腕
            (12, 14),  # 右肩到右肘
            (14, 16),  # 右肘到右手腕
            (11, 23),  # 左肩到左髋
            (12, 24),  # 右肩到右髋
            (23, 24),  # 左右髋
            (23, 25),  # 左髋到左膝
            (25, 27),  # 左膝到左脚踝
            (24, 26),  # 右髋到右膝
            (26, 28),  # 右膝到右脚踝
        ]

        # 绘制骨架连接线
        for start, end in skeleton_connections:
            start_point = (int(keypoints[start].x * image_width), int(keypoints[start].y * image_height))
            end_point = (int(keypoints[end].x * image_width), int(keypoints[end].y * image_height))
            cv2.line(skeleton_image, start_point, end_point, color=(255, 0, 0), thickness=2)
        mp_drawing.draw_landmarks(
            skeleton_image, 
            results.pose_landmarks, 
            mp_pose.POSE_CONNECTIONS,
            mp_drawing.DrawingSpec(color=(0, 255, 0), thickness=7, circle_radius=4),  # 点样式
            mp_drawing.DrawingSpec(color=(255, 0, 0), thickness=7)  # 线样式
        )
        # 保存骨架图像
        file_name, file_extension = os.path.splitext(input_image_path)
        output_image_path = file_name + "_SK" + file_extension
        cv2.imwrite(output_image_path, skeleton_image)
        print(f"骨架图像已保存到 {output_image_path}")
    else:
        # 在图像上显示“不合格”字样
        cv2.putText(
            image, 
            "Picture Unqualified", 
            (50, 50), 
            cv2.FONT_HERSHEY_SIMPLEX, 
            1, 
            (0, 0, 255), 
            2
        )
        # 保存并退出
        print("未检测到人体姿态，图片不合格")


if __name__=="__main__":
    parser=argparse.ArgumentParser(description='SKpoint')
    parser.add_argument("--img_in")
    args=parser.parse_args()
    extract_skeleton_with_validation(args.img_in)
